#!/usr/bin/env bash

# usage:
#     ./imagenet_prepare.sh [imagenet tar file directory] [synset file]

set -e

if [ $# -lt 2 ]; then
    echo "USAGE:"
    echo "    $0 <imagenet directory> <synset file>"
    exit
fi

OUT_DIR="${1}"
SYNSETS_FILE="${2}"
BBOX_DIR="${OUT_DIR}/bounding_boxes/"
CURR_DIR=$(pwd)

if [ ! -d ${OUT_DIR} ]; then
    mkdir -p ${OUT_DIR}
fi
cd "${OUT_DIR}"

BBOX_TARBALL="${OUT_DIR}/ILSVRC2012_bbox_train_v2.tar.gz"
if [ ! -f ${BBOX_TARBALL} ]; then
    echo "${BBOX_TARBALL} DOES NOT EXIST!"
    echo "PLEASE CHECK!"
    exit
fi
if [ ! -d ${BBOX_DIR} ]; then
    mkdir -p ${BBOX_DIR}
else
    rm -rf "${BBOX_DIR}/*"
fi
# Uncompress the BBOX tar ball.
tar xzf "${BBOX_TARBALL}" -C "${BBOX_DIR}"

LABELS_ANNOTATED="${BBOX_DIR}/*"
NUM_XML=$(ls -1 ${LABELS_ANNOTATED} | wc -l)
echo "Identified ${NUM_XML} bounding box annotations."

# Uncompress all images from the ImageNet 2012 validation dataset.
VALIDATION_TARBALL="ILSVRC2012_img_val.tar"
VALIDATION_DIR="${OUT_DIR}/val/"
if [ ! -d ${VALIDATION_DIR} ]; then
    mkdir -p "${VALIDATION_DIR}"
fi
cd "${OUT_DIR}"
if [ ! -f ${VALIDATION_TARBALL} ]; then
    echo "${VALIDATION_TARBALL} DOES NOT EXIST!"
    echo "PLEASE CHECK!"
    exit
fi
tar xf "${VALIDATION_TARBALL}" -C "${VALIDATION_DIR}"

# Prepare the validation data.
${CURR_DIR}/pytorch_val_prep.sh "${VALIDATION_DIR}"

# Uncompress all images from the ImageNet 2012 train dataset.
TRAIN_TARBALL="ILSVRC2012_img_train.tar"
TRAIN_DIR="${OUT_DIR}/train/"
if [ ! -d ${TRAIN_DIR} ]; then
    mkdir -p "${TRAIN_DIR}"
fi
cd "${OUT_DIR}"
if [ ! -f ${TRAIN_TARBALL} ]; then
    echo "${TRAIN_TARBALL} DOES NOT EXIST!"
    echo "PLEASE CHECK!"
    exit
fi

# Un-compress the individual tar-files within the train tar-file.
echo "Uncompressing individual train tar-balls in the training data."

while read  SYNSET; do
    echo "Processing: ${SYNSET}"

    # Create a directory and delete anything there.
    mkdir -p "${TRAIN_DIR}/${SYNSET}"
    rm -rf "${TRAIN_DIR}/${SYNSET}/*"

    # Uncompress into the directory.
    tar xf "${TRAIN_TARBALL}" "${SYNSET}.tar"
    tar xf "${SYNSET}.tar" -C "${TRAIN_DIR}/${SYNSET}/"
    rm -f "${SYNSET}.tar"

    echo "Finished processing: ${SYNSET}"
done < "${CURR_DIR}/${SYNSETS_FILE}"
